"""
WebSocket连接处理
实现实时进度推送
"""
import json
import logging
from typing import Dict, Set
from fastapi import WebSocket, WebSocketDisconnect
from app.config import settings

logger = logging.getLogger(__name__)

class WebSocketManager:
    """WebSocket连接管理器"""

    def __init__(self):
        # 存储活跃连接 {task_id: {connection_id: websocket}}
        self.active_connections: Dict[str, Dict[str, WebSocket]] = {}
        # 存储连接元数据
        self.connection_metadata: Dict[str, Dict] = {}

    async def connect(self, websocket: WebSocket, task_id: str, connection_id: str = None):
        """建立WebSocket连接"""
        await websocket.accept()

        if connection_id is None:
            connection_id = f"conn_{len(self.connection_metadata) + 1}"

        # 添加到活跃连接
        if task_id not in self.active_connections:
            self.active_connections[task_id] = {}

        self.active_connections[task_id][connection_id] = websocket

        # 保存连接元数据
        self.connection_metadata[connection_id] = {
            "task_id": task_id,
            "connected_at": json.dumps({"timestamp": "now"})  # 简化的时间戳
        }

        logger.info(f"WebSocket connected: {connection_id} for task {task_id}")

        # 发送连接确认
        await self.send_to_connection(connection_id, {
            "type": "connected",
            "task_id": task_id,
            "connection_id": connection_id,
            "message": "WebSocket连接已建立"
        })

        return connection_id

    def disconnect(self, connection_id: str):
        """断开WebSocket连接"""
        if connection_id in self.connection_metadata:
            task_id = self.connection_metadata[connection_id]["task_id"]

            # 从活跃连接中移除
            if task_id in self.active_connections:
                self.active_connections[task_id].pop(connection_id, None)
                if not self.active_connections[task_id]:
                    del self.active_connections[task_id]

            # 删除连接元数据
            del self.connection_metadata[connection_id]

            logger.info(f"WebSocket disconnected: {connection_id} for task {task_id}")

    async def send_to_connection(self, connection_id: str, message: dict):
        """向特定连接发送消息"""
        if connection_id in self.connection_metadata:
            task_id = self.connection_metadata[connection_id]["task_id"]
            if (task_id in self.active_connections and
                connection_id in self.active_connections[task_id]):

                websocket = self.active_connections[task_id][connection_id]
                try:
                    await websocket.send_text(json.dumps(message, ensure_ascii=False))
                except Exception as e:
                    logger.error(f"Failed to send message to {connection_id}: {e}")
                    # 清理无效连接
                    self.disconnect(connection_id)

    async def send_to_task(self, task_id: str, message: dict):
        """向任务的所有连接发送消息"""
        if task_id in self.active_connections:
            disconnected_connections = []

            for connection_id, websocket in self.active_connections[task_id].items():
                try:
                    await websocket.send_text(json.dumps(message, ensure_ascii=False))
                except Exception as e:
                    logger.error(f"Failed to send message to connection {connection_id}: {e}")
                    disconnected_connections.append(connection_id)

            # 清理无效连接
            for connection_id in disconnected_connections:
                self.disconnect(connection_id)

    async def broadcast(self, message: dict):
        """向所有连接广播消息"""
        for task_id in list(self.active_connections.keys()):
            await self.send_to_task(task_id, message)

    def get_connection_count(self, task_id: str = None) -> int:
        """获取连接数量"""
        if task_id:
            return len(self.active_connections.get(task_id, {}))
        return sum(len(connections) for connections in self.active_connections.values())

    def get_active_tasks(self) -> Set[str]:
        """获取有活跃连接的任务ID列表"""
        return set(self.active_connections.keys())


# 全局WebSocket管理器实例
websocket_manager = WebSocketManager()